import os
from flask import jsonify
from volcenginesdkarkruntime import Ark


def chat(role_content, user_message):
    # 请确保您已将 API Key 存储在环境变量 ARK_API_KEY 中
    token_my = os.environ.get("ARK_API_KEY")
    # 初始化Ark客户端，从环境变量中读取您的API Key
    client = Ark(
        # 此为默认路径，您可根据业务所在地域进行配置
        base_url="https://ark.cn-beijing.volces.com/api/v3",
        # 从环境变量中获取您的 API Key。此为默认方式，您可根据需要进行修改
        api_key=os.environ.get("ARK_API_KEY"),
    )
    # f""
    try:
        stream = client.chat.completions.create(
            model="doubao-1-5-thinking-pro-250415",
            messages=[
                {"role": "system", "content": role_content},
                {"role": "user", "content": user_message},
            ],
            # 响应内容是否流式返回
            stream=True,
        )
        response_text = ""
        for chunk in stream:
            if not chunk.choices:
                continue
            response_text += chunk.choices[0].delta.content
        return response_text
    except Exception as e:
        return "ERROR"


def chat_with_theme(data, theme, user_ask = None):
    role_content = f"你是宏观经济数据{theme}分析师."
    if not user_ask:
        user_ask = f"请分为三点，每个要点30-50字来描述下面的数据，侧重于最新一期的数据描述与解读,输出不要用markdown格式\n"
    user_message = data.to_csv(sep='\t', na_rep='nan')
    # user_message = "你好"
    response = chat(role_content, user_ask + user_message)
    print('chat response/n', response)
    return response


if __name__ == '__main__':
    import pandas as pd
    role_content = "你是宏观经济数据房地产分析师."
    data = pd.read_excel('meta/test_llm_sample.xlsx', index_col=0)
    
    user_ask = "请分为三点，每个要点30-50字来描述下面的数据，侧重于最新一期的数据描述与解读,输出不要用markdown格式\n"
    user_message = data.to_csv(sep='\t', na_rep='nan')
    # user_message = "你好"
    response = chat(role_content, user_ask + user_message)
    print(response)